#! /usr/bin/env python
# coding=utf-8



# 读取每个OGS 中的paml结果 # 同时生成模拟的文件


mlc = 'mlc'
anc = 'rst'
dat_file = r'my.dat'
ancester_file = r'anc.fa'
aa_dat = r'MCaa.dat'
otherpara = r'para'
new_spe_tree = r'spe_tree'

# 这里有个过滤 就是吧没有找出树的排除
# 产生一个新的目录 重新开始

file_lista = [
'mRNA_cut_ali.fa',
'mRNA_raw_ali.fa',
'pep_cut_ali.fa',
'pep_raw_ali.fa',
'mlc',
'cut_loci',
'rst',

]

total_len = 0
total_gene = 0

aa_order = [ 'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K' ,'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V']

import os

from ete3 import PhyloTree

import shutil
import random





import argparse
import sys



parser = argparse.ArgumentParser(
    description='''读取 第一步aaml的结果 将其放到另外一个文件中去 并为 evolvor 做准备
    这里会过滤部分 没有被aaml跑成功的 序列 这些序列也不会进入后续的支长分析
    用法:
    conver02_read_paml01.py  -i OGS_for_testing -o OGS_for_second_step
    由大天才于2021年11月12日创建于浙江农业大学''')


parser.add_argument('-i',
                help='必须给定,运行完aaml的ogs目录')


parser.add_argument('-o',
                help='必须给定,输出的路径')




args = parser.parse_args()

if not args.i or not args.o:
    parser.print_help()
    sys.exit()



infile = args.i
#outpath = r'OGS_for_second_step'

outpath = args.o

#tree_file = args.t

#orthopath = args.a

#mRNA_path = args.m







try:
	shutil.rmtree(outpath)
except:
	pass

try:
	os.mkdir(outpath)
except:
	pass



for i in os.listdir(infile):
	
	node_nub = None
	ancestor_seq = ''
	tree = ''
	k=None
	alpha = None
	seq_nub = 0
	seq_len = 0
	aa_rate_text = ''
	obj_file = infile+'/'+i+'/'
	aa_dic = { 
	'A':0,
	'R':0,
	'N':0,
	'D':0,
	'C':0,
	'Q':0,
	'E':0,
	'G':0,
	'H':0,
	'I':0,
	'L':0,
	'K':0,
	'M':0,
	'F':0,
	'P':0,
	'S':0,
	'T':0,
	'W':0,
	'Y':0,
	'V':0,
	}

	if os.path.isfile(infile+'/'+i+'/'+mlc):
		# 读取 alpha值 和带有支长的进化树
		with open(infile+'/'+i+'/'+mlc) as fila:
			tree_triger = 0

			for j in fila:
				j = j.strip()
				if tree_triger==4:
					tree = j
					tree_triger = 0

				if tree_triger >=1:
					tree_triger += 1
				if j.find('tree length = ')==0:
					tree_triger = 1
				if j.find('alpha (gamma, K = ') ==0:
					n1 = j.split('alpha (gamma, K = ')[1]
					k = n1.split(')')[0].strip()
					alpha = n1.split('=')[1].strip()
					#print(k,alpha)

				


	else:
		continue
	
	if os.path.isfile(infile+'/'+i+'/'+anc):
		with open(infile+'/'+i+'/'+anc) as fila:
			tree_triger = 0

			for j in fila:

				# 读取root
				j = j.strip()
				
				if tree_triger ==1:
					node_nub = j.split(')')[-1].split(';')[0].strip()
				if tree_triger >=1:
					tree_triger += 1
				if j.find('tree with node labels for Rod Page') ==0:
					tree_triger = 1
				if node_nub!=None:
					
					if j.find('node #'+node_nub)==0:

						ancestor_seq = j[len('node #'+node_nub)+1:].replace(' ','')
						# 计算各个氨基酸比例
						for aa in ancestor_seq:
							aa_dic[aa] += 1
						aa_dic ={aa:aa_dic[aa]/len(ancestor_seq) for aa in aa_dic}



		#print(aa_dic)
	else:
		continue
	if tree=='':
		print('OGS '+i+' 没能获得带支长的树' )
		continue



	ptree = PhyloTree(tree)

	seq_nub = len(ptree.get_leaf_names())
	seq_len = len(ancestor_seq)
	aa_rate_text += ' '.join([str(aa_dic[j]) for j in aa_order[:10]])+'\n'
	aa_rate_text += ' '.join([str(aa_dic[j]) for j in aa_order[10:]])

	#print(aa_rate_text)
	total_len += seq_len
	total_gene += 1

	evolver_temp = '''
	0        * 0: paml format (mc.paml); 1:paup format (mc.nex)
    %s       * random number seed (odd number)
    
    %s %s 1   * <# seqs>  <# sites>  <# replicates>
    
    -1         * <tree length, use -1 if tree below has absolute branch lengths>
    
    %s
    
    %s %s        * <alpha; see notes below>  <#categories for discrete gamma>
    3 %s * <model> [aa substitution rate file, need only if model=2 or 3]
    
    %s
    
     A R N D C Q E G H I
     L K M F P S T W Y V
    
    // end of file
    
    =============================================================================
    Notes for using the option in evolver to simulate amino acid sequences. 
    Change values of parameters, but do not delete them.  It is o.k. to add 
    empty lines, but do not break down the same line into two or more lines.
    
      model = 3 (poisson), 1 (proportional), 2 (empirical), 3 (empirical_F)
      Use 0 for alpha to have the same rate for all sites.
      Use 0 for <#categories for discrete gamma> to use the continuous gamma
      <aa substitution rate file> can be dayhoff.dat, jones.dat, and so on.
      <aa frequencies> have to be in the right order, as indicated.
    =================!! Check screen output carefully!! =====================
	''' %(int(random.random()*1000),seq_nub, seq_len, tree, alpha, k, dat_file, aa_rate_text)

	# 从新 建立文件
	#print(i)
	os.mkdir(outpath+'/'+i)
	target_path = outpath+'/'+i+'/'
	with open(target_path+ancester_file,'w' ) as fila:
		fila.write('>anc\n'+ancestor_seq+'\n')
	with open(target_path+aa_dat,'w' ) as fila:
		fila.write(evolver_temp)

	with open(target_path+otherpara,'w' ) as fila:
		fila.write('k = '+k+'\n')
		fila.write('alpha = '+alpha+'\n')

	with open(target_path+new_spe_tree,'w' ) as fila:
		fila.write(ptree.write(format=1))

	shutil.copy2(obj_file +dat_file, target_path)
	
	for z in file_lista:
		shutil.copy2(obj_file +z, target_path)

print('过滤得到了 %s 个OGS  %s 个aa' %(total_gene, total_len))